package org.example.apitest.sink;

import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.functions.RuntimeContext;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.connectors.elasticsearch.ElasticsearchSinkFunction;
import org.apache.flink.streaming.connectors.elasticsearch.RequestIndexer;
import org.apache.flink.streaming.connectors.elasticsearch7.ElasticsearchSink;
import org.apache.http.HttpHost;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.client.Requests;
import org.example.apitest.beans.SensorReading;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Map;

public class SinkTest3ES {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        // 从文件中读取数据
        DataStream<String> inputStream = env.readTextFile("src/main/resources/sensor.txt");

        DataStream<SensorReading> dataStream = inputStream.map((MapFunction<String, SensorReading>) SensorReading::new);

        dataStream.addSink(new ElasticsearchSink.Builder<SensorReading>(new ArrayList<HttpHost>() {{
            add(new HttpHost("localhost", 9200, "http"));
        }}, new ElasticsearchSinkFunction<SensorReading>() {
            @Override
            public void process(SensorReading sensorReading, RuntimeContext runtimeContext, RequestIndexer requestIndexer) {
                // 定义写入数据的source
                Map<String, String> dataSource = new HashMap<String, String>() {{
                    put("id", sensorReading.getId());
                    put("temp", sensorReading.getTemperature().toString());
                    put("ts", sensorReading.getTimestamp().toString());
                }};

                // 创建请求,作为向ES发起的写入命令
                IndexRequest indexRequest = Requests.indexRequest().index("sensor").type("readingData").source(dataSource);

                // 用index发送请求
                requestIndexer.add(indexRequest);
            }
        }).build());
        env.execute();
    }
}
